import matplotlib.pyplot as plt

# 给定假设数据集
x_data = [1.0, 2.0, 3.0, 4.0]
y_data = [2.0, 4.0, 6.0, 8.0]

# 初始权重
w = 1.0


def forward(x):
    """
    定义模型：y' = w * x

    Args:
        x:输入数据，样本
    """
    return x * w


def cost(xs, ys):
    """
    计算损失（或者代价）

    Args:
        xs:输入数据集
        ys:输入数据集的标签
    """
    # 损失（或代价）
    cost = 0;
    for x, y in zip(xs, ys):
        y_pred = forward(x)
        cost += (y_pred - y) ** 2
    return cost / len(xs)


def gradient(xs, ys):
    """
    计算梯度

    Args:
        xs:输入数据集
        ys:输入数据集的标签
    """
    grad = 0;
    for x, y in zip(xs, ys):
        grad += 2 * x * (x * w - y)
    return grad / len(xs)


# 初始计算
print('Predict (before training)', 4, forward(4))

# 存放每一轮损失的值
cost_list = [1.0]
# 迭代训练
for epoch in range(100):
    # 计算损失
    cost_val = cost(x_data, y_data)
    cost_list.append(cost_val)
    # 计算梯度，这里使从新遍历数据集计算的，比较费时间
    grad_val = gradient(x_data, y_data)
    # 更新权重
    w -= 0.01 * grad_val
    # 输出每一轮的结果
    print(('Epoch:', epoch, 'w=', w, 'loss=', cost_val))

print('Predict (before training)', 4, forward(4))
# 绘图，传入数据x、y、[fmt]，[fmt]是线条的格式，包括颜色等
plt.plot(range(101), cost_list)
# 纵、横坐标
plt.ylabel('Cost')
plt.xlabel('Epoch')
# 图名
plt.title("Gradient Descent")
# 展示
plt.show()